import matplotlib.pyplot as plt
from scipy.spatial.transform import Rotation as R


def visualize_sensors(quaternions, positions):

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    max_pos = max(max(p) for p in positions) if positions else 1
    ax.set_xlim([-max_pos - 0.5, max_pos + 0.5])
    ax.set_ylim([-max_pos - 0.5, max_pos + 0.5])
    ax.set_zlim([-max_pos - 0.5, max_pos + 0.5])

    # Loop all sensors
    for i, (q, pos) in enumerate(zip(quaternions, positions)):

        rot = R.from_quat([q[3], q[0], q[1], q[2]])
        rot_mat = rot.as_matrix()

        ax.quiver(*pos, *rot_mat[:, 0], color='r', length=0.3, normalize=True)
        ax.quiver(*pos, *rot_mat[:, 1], color='g', length=0.3, normalize=True)
        ax.quiver(*pos, *rot_mat[:, 2], color='b', length=0.3, normalize=True)
        ax.text(*pos, f'Sensor {i+1}', fontsize=8)

    ax.set_xlabel('Global X')
    ax.set_ylabel('Global Y')
    ax.set_zlabel('Global Z')
    plt.title('Sensor Coordinate Systems')
    plt.show()


if __name__ == "__main__":

    # Rotation
    sensors_quat = [
        [0.7071, -0.7071, 0, 0],  # Sensor 1
        [0, 0, 0, 1],  # Sensor 2（单位四元数）
        [0.5, 0.5, 0.5, 0.5]  # Sensor 3（示例四元数）
    ]

    # Pos: [X, Y, Z]
    sensor_positions = [[0, 0, 0], [1, 0, 0], [0, 1, 0]]

    visualize_sensors(sensors_quat, sensor_positions)
